import json
import os
import re
import sys
from itertools import combinations, repeat, takewhile

import bpy
from bpy.types import NodeTree, Node, NodeSocket

import igraph as ig

'''
输出json格式内容
'''

def dumps(obj):
    """Outputs json with formatting edits + object handling."""
    return json.dumps(obj, indent=4, sort_keys=True, cls=CustomEncoder)


class CustomEncoder(json.JSONEncoder):

    def encode(self, obj):
        """Fired for every object."""
        s = super(CustomEncoder, self).encode(obj)
        # If uncompressed, postprocess for formatting
        if len(s.splitlines()) > 1:
            s = self.postprocess(s)
        return s

    def postprocess(self, json_string):
        """Displays each entry on its own line."""
        is_compressing, is_hash, compressed, spaces = False, False, [], 0
        for row in json_string.split("\n"):
            if is_compressing:
                if (row[:spaces + 5] == " " * (spaces + 4) +
                        ("\"" if is_hash else "{")):
                    compressed.append(row.rstrip())
                elif (len(row) > spaces and row[:spaces] == " " * spaces and
                        re.match("[\]\}],?", row[spaces:].rstrip())):
                    compressed.append(row.rstrip())
                    is_compressing = False
                else:
                    compressed[-1] += " " + row.strip()
            else:
                compressed.append(row.rstrip())
                if any(a in row for a in ["edges", "nodes"]):
                    # Fix to handle issues that arise with empty lists
                    if "[]" in row:
                        continue
                    spaces = sum(1 for _ in takewhile(str.isspace, row))
                    is_compressing, is_hash = True, "{" in row
        return "\n".join(compressed)


# blender add-on部分——tool shelf中的Dynamic display

class memoryDynamicPanel(bpy.types.Panel):
    bl_idname = 'memory.DynamicDisplay'
    bl_label = 'Dynamic Display'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'TOOLS'
    bl_category = 'Memory'

    def draw(self, context):
        col = self.layout.column(align=True)
        col.label(text="Dynamic control:")
        row = col.row(align=True)
        row.operator("memory.prev", text="Prev")
        row.operator("memory.next", text="Next")

    @classmethod
    def register(cls):
        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))

#
#    控制前进，动态显示
#
class MEMORY_OT_NextButton(bpy.types.Operator):
    bl_idname = "memory.next"
    bl_label = "Next Memory"
    
    def execute(self, context):
        return{'FINISHED'}

#
#   控制后退
#
class MEMORY_OT_PrevButton(bpy.types.Operator):
    bl_idname = "memory.prev"
    bl_label = "previous Memory"
    
    def execute(self, context):
        return{'FINISHED'}


class memoryColorPanel(bpy.types.Panel):
    bl_idname = 'memory.ColorInfo'
    bl_label = 'Color Info'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'TOOLS'
    bl_category = 'Memory'

    def draw(self, context):
        purple_ball = icon_collections["purple_ball"]["purple_ball"]
        gray_ball = icon_collections["gray_ball"]["gray_ball"]
        green_ball = icon_collections["green_ball"]["green_ball"]
        yellow_ball = icon_collections["yellow_ball"]["yellow_ball"]
        blue_ball = icon_collections["blue_ball"]["blue_ball"]
        red_ball = icon_collections["red_ball"]["red_ball"]
        white_ball = icon_collections["white_ball"]["white_ball"]

        col = self.layout.column(align=True)
        col.label(text=": Heaps", icon_value=purple_ball.icon_id)
        col.label(text=": DLLs", icon_value=gray_ball.icon_id)
        col.label(text=": Stacks", icon_value=green_ball.icon_id)
        col.label(text=": Mapped Files", icon_value=yellow_ball.icon_id)
        col.label(text=": root node", icon_value=blue_ball.icon_id)
        col.label(text=": new or changed nodes", icon_value=red_ball.icon_id)
        col.label(text=": normal VAD nodes", icon_value=white_ball.icon_id)

    @classmethod
    def register(cls):
        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))


class memoryInfoPanel(bpy.types.Panel):
    bl_idname = 'memory.VADInfo'
    bl_label = 'VAD Info'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'TOOLS'
    bl_category = 'Memory'

    def draw(self, context):
        if context.active_node:
            node_info = vad_tree["nodes"][context.active_node.name]
            col = self.layout.column(align=True)
            col.label(text="pid: "+str(node_info["pid"]))
            col.label(text="vad: "+hex(node_info["vad"]))
            col.label(text="start: "+hex(node_info["start"]))
            col.label(text="end: "+hex(node_info["end"]))
            col.label(text="VadTag: "+str(node_info["VadTag"]))
            col.label(text="VadType: "+str(node_info["VadType"]))
            col.label(icon="IPO_SINE")
            col.label(text="flags: "+str(node_info["flags"]))
            col.label(text="protection: "+str(node_info["protection"]))
            col.label(icon="IPO_QUAD")
            col.label(text="ControlArea: "+hex(node_info["ControlArea"]))
            col.label(text="segment: "+hex(node_info["segment"]))
            col.label(icon="IPO_CUBIC")
            col.label(text="NumberOfSectionReferences: "+str(node_info["NumberOfSectionReferences"]))
            col.label(text="NumberOfPfnReferences: "+str(node_info["NumberOfPfnReferences"]))
            col.label(text="NumberOfMappedViews: "+str(node_info["NumberOfMappedViews"]))
            col.label(text="NumberOfUserReferences: "+str(node_info["NumberOfUserReferences"]))
            col.label(icon="IPO_QUART")
            col.label(text="ControlFlags: "+str(node_info["ControlFlags"]))
            col.label(text="FileObject: "+hex(node_info["FileObject"]))
            col.label(text="FileName: "+str(node_info["FileName"]))
            col.label(icon="IPO_QUINT")
            col.label(text="FirstprototypePTE: "+hex(node_info["FirstprototypePTE"]))
            col.label(text="LastcontiguousPTE: "+hex(node_info["LastcontiguousPTE"]))
            col.label(text="Flags2: "+str(node_info["Flags2"]))
        

    @classmethod
    def register(cls):
        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))


class VadCustomTree(NodeTree):
    '''A vad node tree type that will show up in the node editor header'''
    bl_idname = 'VadTreeType'
    bl_label = 'Vad Node Tree'
    bl_icon = 'NODETREE'

# Vad socket type
class VadCustomSocket(NodeSocket):
    '''Vad node socket type'''
    bl_idname = 'VadSocketType'
    bl_label = 'Vad Node Socket'

    def draw(self, context, layout, node, text):
        layout.label(text)

    def draw_color(self, context, node):
        return (1.0, 0.4, 0.216, 0.5)


class VadCustomTreeNode:
    @classmethod
    def poll(cls, ntree):
        return ntree.bl_idname == 'VadTreeType'


class OBJECT_OT_VADButton(bpy.types.Operator):
    bl_idname = "vad.address"
    bl_label = "Vad Address"
    
    def execute(self, context):
        return{'FINISHED'}


class VadCustomNode(Node, VadCustomTreeNode):
    '''A custom node'''
    bl_idname = 'VadNodeType'
    bl_label = 'Vad Node'
    bl_icon = 'SOUND'

    VAD = bpy.props.StringProperty(default="0x70238a30 - 0x7ffd2fff")

    def init(self, context):
        self.inputs.new('VadSocketType', " ")

        self.outputs.new('NodeSocketColor', " ")
    
    def copy(self, node):
        pass
    
    def free(self):
        print("Removing node ", self, ", Goodbye!")
    
    def draw_buttons(self, context, layout):
        tmp = self.VAD
        layout.label(text=str(tmp))
    
    def draw_label(self):
        return "vad node"

import nodeitems_utils
from nodeitems_utils import NodeCategory, NodeItem

class VadNodeCategory(NodeCategory):
    @classmethod
    def poll(cls, context):
        return context.space_data.tree_type == 'VadTreeType'

node_categories = [
    # identifier, label, items list
    VadNodeCategory("SOMENODES", "Some Nodes", items=[
        # our basic node
        NodeItem("VadNodeType"),
        ]),
    ]

# icon collection
icon_collections = {}

def register():
    import bpy.utils.previews
    icon = bpy.utils.previews.new()
    icons_dir = "../../icons"
    for ball in os.listdir(icons_dir):
        icon.load(ball[:-4], os.path.join(icons_dir, ball), 'IMAGE')
        icon_collections[ball[:-4]] = icon
    bpy.utils.register_class(memoryColorPanel)
    bpy.utils.register_class(memoryDynamicPanel)
    bpy.utils.register_class(MEMORY_OT_NextButton)
    bpy.utils.register_class(MEMORY_OT_PrevButton)
    bpy.utils.register_class(memoryInfoPanel)
    bpy.utils.register_class(OBJECT_OT_VADButton)
    bpy.utils.register_class(VadCustomTree)
    bpy.utils.register_class(VadCustomSocket)
    bpy.utils.register_class(VadCustomNode)

    nodeitems_utils.register_node_categories("CUSTOM_NODES", node_categories)


def unregister():
    nodeitems_utils.unregister_node_categories("CUSTOM_NODES")
    bpy.utils.unregister_class(OBJECT_OT_VADButton)
    bpy.utils.unregister_class(VadCustomTree)
    bpy.utils.unregister_class(VadCustomSocket)
    bpy.utils.unregister_class(VadCustomNode)
    bpy.utils.unregister_class(memoryDynamicPanel)
    bpy.utils.unregister_class(memoryInfoPanel)
    bpy.utils.unregister_class(MEMORY_OT_NextButton)
    bpy.utils.unregister_class(MEMORY_OT_PrevButton)
    for icon in icon_collections.values():
        bpy.utils.previews.remove(icon)
    icon_collections.clear()
    bpy.utils.unregister_class(memoryColorPanel)


# 颜色RGB指定
colors = { "red": (1.0, 0.411, 0.503), "yellow": (0.981, 1.0, 0.272), "green": (0.563, 1.0, 0.541),
           "blue": (0.241, 0.861, 1.0), "gray": (0.4, 0.4, 0.4), "white": (1.0, 1.0, 1.0), 
           "purple": (0.58, 0.579, 1.0)}


if __name__ == "__main__":
    pid = "408"
    register()
    try:
        with open(r"json/vad0.json") as in_file:
            net = json.load(in_file)
    except IOError:
        with open(sys.argv[-1]) as in_file:
            net = json.load(in_file)

    labels = list(net["nodes"][pid].keys())
    edges = net["edges"][pid]
    Edges = []
    def generate_edges(*args, **kw):
        if kw['source'] in args:
            i = args.index(kw["source"])
        if kw['target'] in args:
            j = args.index(kw["target"])
        tup = (i, j)
        Edges.append(tup)
    for k in range(len(edges)):
        generate_edges(*labels, **edges[k])
    rt = []
    rt.append(Edges[0][0])
    G = ig.Graph(Edges, directed = False)
    layt = G.layout_reingold_tilford(mode="in",root=rt)
    layt.center()
    master_nodes = net["nodes"][pid]
    for i in range(len(labels)):
        master_nodes[labels[i]]["location"] = [j*200 for j in layt[i]]
        tmp = master_nodes[labels[i]]["location"][0]
        master_nodes[labels[i]]["location"][0] = master_nodes[labels[i]]["location"][1]
        master_nodes[labels[i]]["location"][1] = tmp

    json_str = dumps({"edges": edges, "nodes": master_nodes})
    vad_tree = json.loads(json_str)

    ### blender operation ###
    bpy.context.scene.use_nodes = True
    tree = bpy.context.scene.node_tree

    # clear default nodes
    for node in tree.nodes:
        tree.nodes.remove(node)

    # create input vad node
    for vad_label in labels:
        vad_node = tree.nodes.new(type='VadNodeType')
        vad_node.location = tuple(vad_tree["nodes"][vad_label]["location"])
        vad_node.name = vad_label
        vad_node.label = vad_label
        vad_node.VAD = vad_tree["nodes"][vad_label]["address"]
        vad_node.use_custom_color = True
        vad_node.color = colors[vad_tree["nodes"][vad_label]["color"]]

    links = tree.links
    for e in vad_tree["edges"]:
        source = tree.nodes[e["source"]]
        target = tree.nodes[e["target"]]
        links.new(source.outputs[0], target.inputs[0])
